from driving_gym.environment.agent.done.base_done import BaseDone
from driving_gym.environment.scenario.carla_scenario import CarlaScenario
from driving_gym.simulation.adapter.carla.carla_adapter import CarlaAdapter
from driving_gym.simulation.adapter.carla.carla_config import CarlaActorConfig


class CarlaDefaultDone(BaseDone):
    adapter: CarlaAdapter

    def __init__(self, config: dict, adapter=None):
        super().__init__(config, adapter)

        self.reach_threshold = self.params.get("reach_threshold", 5)

    def reset(self):
        self.steps = 0

    def get_done(self, snapshot, scenario: CarlaScenario):
        transform = snapshot.data[self.actor_id]["transform"]
        goal = CarlaActorConfig.parse_location(
            scenario.get_task(self.actor_id)["end"], self.adapter.map
        )
        if scenario.enable_planner:
            planner = scenario.planners[self.actor_id]
        else:
            planner = None

        # Reach
        if not scenario.endless:
            if transform.location.distance(goal.location) < self.reach_threshold:
                return True, "REACH_GOAL"
            if planner and planner.get_distance_to_end() < self.reach_threshold:
                return True, "REACH_GOAL"

        # Collision
        if scenario.attach_default_sensors:
            collision = snapshot.data[self.actor_id]["sensors"]["collision"]
            if collision[1] > 0:
                return True, "COLLIDE_VEHICLE"
            if collision[2] > 0:
                return True, "COLLIDE_PEDESTRIAN"
            if collision[3] > 0:
                return True, "COLLIDE_OTHER"

        return False, "NOT_DONE"
